package com.arch.flink.window.watermark;

import com.alibaba.fastjson.JSONObject;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * @author pizhihui
 * @date 2024-03-11
 */
public class Demo3New {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<String> source = env.socketTextStream("127.0.0.1", 9999);

        SingleOutputStreamOperator<String> outputStreamOperator =
                source.assignTimestampsAndWatermarks(WatermarkStrategy
                                .<String>forBoundedOutOfOrderness(Duration.ofSeconds(3))
                                .withTimestampAssigner(new SerializableTimestampAssigner<String>() {
                                    @Override
                                    public long extractTimestamp(String t, long l) {
                                        String[] s = t.split(" ");
                                        System.out.println("event time: " + s[0]);
                                        return Long.parseLong(s[0]);
                                    }
                                }))
                        .map(new MapFunction<String, Tuple2<String, Integer>>() {
                            @Override
                            public Tuple2<String, Integer> map(String s) throws Exception {
                                return Tuple2.of(s.split(" ")[1], 1);
                            }
                        })
                        .keyBy(s -> s.f0)
                        .window(TumblingEventTimeWindows.of(Time.seconds(10)))
                        .allowedLateness(Time.seconds(3))
                        .process(new ProcessWindowFunction<Tuple2<String, Integer>, String, String, TimeWindow>() {
                            @Override
                            public void process(String s,
                                                ProcessWindowFunction<Tuple2<String, Integer>,
                                                        String, String, TimeWindow>.Context context,
                                                Iterable<Tuple2<String, Integer>> iterable,
                                                Collector<String> collector) throws Exception {
                                int sum = 0;
                                for (Tuple2<String, Integer> item : iterable) {
                                    sum += item.f1;
                                }
                                JSONObject jsonObject = new JSONObject();
                                jsonObject.put("windowEnd", context.window().getEnd());
                                jsonObject.put("waterMark", context.currentWatermark());
                                jsonObject.put("key", s);
                                jsonObject.put("sum", sum);
                                collector.collect(jsonObject.toJSONString());
                            }
                        });

        outputStreamOperator.print("out");
        env.execute();


    }

}
